import json
import secrets
import string

import jwt
from flask_sock import Sock
from flask import request, current_app
from app import create_app, mongo
from app.models import Prompt,UserTopicAssociation,Message
from app.ai import get_assistant_response



app = create_app()
sock = Sock(app)

def generate_random_topic_id(length=16):
    characters = string.ascii_letters + string.digits
    return ''.join(secrets.choice(characters) for _ in range(length))


@sock.route('/chat')
def chat(ws):
    authorization = request.args.get('Authorization')
    user_id = request.args.get('userId')

    ##userid和token的是否一样

    #2.topicid

    topic_id = request.args.get('topicId')
    ##鉴定权限,如果第一次xxxtopicID生成 第二次xx
    ##退出保存到mogodb 第二次有topic 直接读取mogodb聊天记录

    ##无论怎么样 一样修改到数据库

    if not authorization or ' ' not in authorization:
        ws.close()
        return

    parts = authorization.split()
    if len(parts) == 2 and parts[0].lower() == 'bearer':
        token = parts[1]


    try:
        data = jwt.decode(token, current_app.config['SECRET_KEY'], algorithms=['HS256'])
        # 从 token 中获取 user_id，然后比较用户id
        # 校验 user_id 是否匹配 token 中的 user_id
        if user_id != str(data.get('user_id')):
            ws.close()
            return
    except jwt.ExpiredSignatureError:
        # token 过期
        ws.close()
        return
    except jwt.InvalidTokenError:
        # token 无效
        ws.close()
        return

        # 如果 topic_id 为空或不存在，则生成一个 16 位随机字符串
    history_message=[]
    user_topic = None
    if not topic_id:
        topic_id = generate_random_topic_id()
        ##第一次聊天
        ##创建topic
        user_topic = UserTopicAssociation(user_id=user_id, topic_id= topic_id,title="新的聊天")
        user_topic.save_to_db()
        mongo_message = {"_id": topic_id, "message": history_message}
        mongo.db.topics.insert_one(mongo_message)
    else:
        print(topic_id)
        user_topic = UserTopicAssociation.query.filter(UserTopicAssociation.topic_id==topic_id).first()

        ##查找出聊天记录 继续聊天
        try:
            mongo_message_find = mongo.db.topics.find_one({"_id": topic_id})
            history_message =mongo_message_find.get('message')

        except Exception as e:
            ws.close()
    #查找topic


    try:



        while True:
            message = ws.receive()
            try:

                input_message = json.loads(message)

                status = input_message['status']

                ##根据status判断时候需要prompt参数
                ###需要加入 既需要加入背景

                if status == 0:



                    promptId = input_message['promptId']

                    ###需要加入参数
                    prompt = Prompt.query.filter(Prompt.id == promptId).first()

                    types = prompt.prompt_name

                    title = input_message['question']
                    prompt_content = prompt.prompt

                    user_topic.types = types
                    user_topic.title = title
                    user_topic.update_from_db()

                    message_prompt = Message(role="user",content=prompt_content)
                    history_message.append(message_prompt.to_dict())

                    ##新聊天 设置参数

                result = get_assistant_response(input_content=input_message['question'] ,conversation_history=history_message)

                if result is not None:
                    history_message.append(Message(role="system",content=result).to_dict())

                ##发送结果
                ws.send(result)

            except json.JSONDecodeError:

                print("Received message is not valid JSON")
                continue
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:

        mongo.db.topics.find_one_and_update(
            {"_id": topic_id},  # 直接使用字符串 ID
            {"$set": {"message": history_message}},
            return_document=False
        )
        ## 更新数据库
        user_topic.update_from_db()

        print("ws 连接已断开")

@app.before_request
def before():
    url = request.path  # 当前请求的URL
    # passUrl = ["/login", "/user/register", "/verificationCode"]
    # if url in passUrl:
    #     pass

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)